"""
DyTyrant - Distributed Tokyo Tyrant in python

Author : Alban Mouton (alban.mouton@gmail.com)

DyTyrant makes it possible to use Tokyo Tyrant through PyTyrant in cluster mode.
It is client side, meaning it doesn't need any additional server nor modifications of Tokyo Tyrant server.

DyTyrant is a very light piece of code for python developers, not a full featured project for DB admins.

Usage example:

Prepare Tokyo Tyrant servers:

user@node1> ttserver -port 2010 foo.tch
user@node1> ttserver -port 2011 bar.tch

user@node2> ttserver -port 2010 foo.tch
user@node2> ttserver -port 2011 bar.tch

test.py:

import dytyrant

# "example.pckl" should be shared by all clients connecting to the same cluster it will be generated by the first client that will try to use it
# Number of replicates per key is set to 2
# If the parameters are incompatibles with the ones of a previously existing file, DyTyrant will attempt to convert its configuration
dyconf = dytyrant.DyConf("example.pckl",hostnames=['node1','node2'],replicas=2,rebalance=True)

# Several dytyrant instances can share a common DyConf instance
db1 = dytyrant.DyTyrant(2010,dyconf)
db2 = dytyrant.DyTyrant(2011,dyconf)

# You can use DyTyrant instances almost as dictionaries
db1["foo"] = "bar"
print db1["foo"]
del db1["foo"]
etc...

# You can repair DyTyrant instances (check position of keys on the cluster and correct if necessary)
db1.repair()
# then close it
db1.close()


TODO:
rebalance() with different set of hosts
Doc for conf_from_db()
dy_mem : allow for reparation on memorized operations (instead of whole tables)
"""

import pytyrant, pickle, random, socket, hashlib, os, logging, threading, anydbm, time

class DyTyrant(dict):
	"""
	Main class of DyTyrant
	"""
	def __init__(self,port,dyconf):
		self.port = port
		self.dyconf = dyconf
		#self.mem = DyMem("/tmp/dymem_"+str(time.time()))
		#self.mem.start()
		self.localhost = socket.gethostname()
		
		# build all pytyrant instances
		self.pytyrants = {}
		self.fail_hosts = set([])
		for host in self.dyconf.hostnames:
			try:
				self.pytyrants[host] = pytyrant.PyTyrant.open(host, port)
			except:
				logging.warning("Connection with host %s failed on port %s" % (host,port))
				self.fail_hosts.add(host)

		# if dyconf had to rebalance, repair
		if self.dyconf.needs_repair:
			self.repair()

	def close(self):
		"""
		Close all pytyrant instance
		"""
		for host in self.pytyrants:
			try:
				self.pytyrants[host].close()
			except:
				continue

	def get_hosts(self,key=None):
		"""
		Get the hosts available for a key (or all hosts available)

		If localhost is available, use it first
		"""
		if key:
			hosts = set(self.dyconf.get_hosts(key)) - self.fail_hosts
		else:
			hosts = set(self.dyconf.hostnames) - self.fail_hosts
		if len(hosts) == 0:
			raise Exception("Too many hosts failed, no host available for key %s" % key)
		
		if self.localhost in hosts:
			lhosts = [self.localhost]
			lhosts.extend(list(hosts - set([self.localhost])))
		else:
			lhosts = list(hosts)
		return lhosts

	def repair(self):
		"""
		Scan all databases and correct keys partitioning if needed
		"""
		mem_keys = set([])
		# scan each db
		for current_host in self.get_hosts():
			db = self.pytyrants[current_host]
			to_delete = []
			for key in db.iterkeys():
				# memorize ckecked keys
				if key in mem_keys:
					continue
				mem_keys.add(key)

				# if the key shouldn't be on this host, remember to delete it later
				hosts = self.get_hosts(key)
				if current_host not in hosts:
					to_delete.append(key)
				else:
					hosts.remove(current_host)

				# check if the key is present on all its replicate hosts
				for host in hosts:
					try:
						if key not in self.pytyrants[host]:
							self.pytyrants[host][key] = db[key]
					except KeyboardInterrupt,e:
						raise e
					except Exception,e:
						logging.warning("PUT failed for host %s" % host)
						self.fail_hosts.add(host)

				# check if the key is absent on all other hosts
				not_hosts = self.get_hosts();not_hosts.remove(current_host)
				not_hosts = set(not_hosts) - set(hosts)
				for host in not_hosts:
					try:
						del self.pytyrants[host][key]
					except KeyboardInterrupt,e:
						raise e
					except:
						pass
			# delete key that had to wait after the end of iteration
			for key in to_delete:
				del db[key]

	def repair_from_mem(self):
		"""
		Try to repair databases for memorized keys
		"""
		for host in self.dyconf.hostnames:
			if host in self.fail_hosts:
				try:
					self.pytyrants[host] = pytyrant.PyTyrant.open(host,self.port)
				except:
					continue
#			for key in self.mem.iter(host):
#				try:




	def conf_from_db(self):
		"""
		Update the dyconf to fit with existing key distribution
		"""
		sample_keys = {}
		for host in self.get_hosts():
			for i,key in enumerate(self.pytyrants[host]):
				sample_keys[key] = sample_keys.get(key,[])
				sample_keys[key].append(host)
				if i>1000:
					break

		self.dyconf.build_from_sample(sample_keys)
		self.dyconf.dump()

	def iterkeys(self):
		"""
		DyTyrant iterators are not very good, they have a memory that can use a lot of space
		"""
		mem_keys = set([])
		for host in self.get_hosts():
			for key in self.pytyrants[host].iterkeys():
				if key in mem_keys:
					continue
				mem_keys.add(key)
				yield key

	def iteritems(self):
		"""
		DyTyrant iterators are not very good, they have a memory that can use a lot of space
		"""
		mem_keys = set([])
		for host in self.get_hosts():
			for key in self.pytyrants[host].iterkeys():
				if key in mem_keys:
					continue
				mem_keys.add(key)
				yield (key,db[value])

	def itervalues(self):
		"""
		DyTyrant iterators are not very good, they have a memory that can use a lot of space
		"""
		mem_keys = set([])
		for host in self.get_hosts():
			for key in self.pytyrants[host].iterkeys():
				if key in mem_keys:
					continue
				mem_keys.add(key)
				yield db[key]

	def has_key(self,key):
		"""
		Call __contains__()
		"""
		return self.__contains__(key)

	def __iter__(self):
		"""
		DyTyrant iterators are not very good, they have a memory that can use a lot of space
		"""
		for key in self.iterkeys():
			yield key

	def __getitem__(self,key):
		for host in self.get_hosts(key):
			try:
				return self.pytyrants[host][key]
			except KeyError:
				logging.warning("GET failed for host %s on port %s (KeyError)" % (host,self.port))
			except KeyboardInterrupt,e:
				raise e
			except Exception,e:
				logging.warning("GET failed for host %s on port %s (%s : %s)" % (host,self.port,e.__class__,e.message))
				self.fail_hosts.add(host)
		raise KeyError(key)

	def __setitem__(self,key,value):
		for host in self.get_hosts(key):
			try:
				self.pytyrants[host][key] = value
			except KeyboardInterrupt,e:
				raise e
			except Exception,e:
				logging.warning("PUT failed for host %s on port %s (%s : %s)" % (host,self.port,e.__class__,e.message))
				self.fail_hosts.add(host)
	
	def __delitem__(self,key):
		for host in self.get_hosts(key):
			try:
				del self.pytyrants[host][key]
			except KeyError:
				logging.warning("DEL failed for host %s on port %s (KeyError)" % (host,self.port))
			except KeyboardInterrupt,e:
				raise e
			except Exception,e:
				logging.warning("DEL failed for host %s on port %s (%s : %s)" % (host,self.port,e.__class__,e.message))
				self.fail_hosts.add(host)

	def __contains__(self,key):
		for host in self.get_hosts(key):
			try:
				if key in self.pytyrants[host]:
					return True
			except KeyboardInterrupt,e:
				raise e
			except Exception,e:
				logging.warning("IN failed for host %s on port %s (%s : %s)" % (host,self.port,e.__class__,e.message))
				self.fail_hosts.add(host)
		return False

	def __len__(self):
		length = 0
		for host in self.get_hosts():
			try:
				length += len(self.pytyrants[host])
			except KeyboardInterrupt,e:
				raise e
			except Exception,e:
				logging.warning("LEN failed for host %s on port %s (%s : %s)" % (host,self.port,e.__class__,e.message))
				self.fail_hosts.add(host)
		return length / self.dyconf.replicas

	def __getattr__(self,name):
		raise NotImplementedError("Function %s is not implemented in DyTyrant" % name)


class DyMem(threading.Thread):
	"""
	Memorize operations that didn't work, and try to apply them regularly
	"""

	def run(self):
		"""
		Execution method
		Atempt regularly to apply memorized operations
		"""
		self.db = anydbm.open(file,'c')
		try:
			while not stop:
				time.sleep(600)
				self.dytyrant.repair_from_mem()
		except Exception,e:
			self.db.close()
			raise e
	
	def get_pos(self,host,op):
		"""
		Find the position in the db to memorize a key for this host (used by mem())
		"""
		i = 0
		while host+"_"+str(i) in self.db:
			i+=1
		return host+"_"+str(i)

	def mem(self,host,key):
		"""
		Remember to check the value of a key on a host
		"""
		self.db[self.get_pos(host)] = key

	def iter(self,host):
		"""
		Iterate on the list of keys that should be set for a host
		"""
		i = 0
		while host+"_"+str(i) in self.db:
			yield self.db[host+"_"+str(i)]
			i += 1

	def erase(self,host):
		"""
		Iterate on the list of keys that should be deleted on a host
		"""
		i = 0
		while host+"_"+str(i) in self.db:
			del self.db[host+"_"+str(i)]
			i += 1
		
		

class DyConf:
	"""
	Configuration and hash functions for DyTyrant
	"""
	def __init__(self,file,hostnames=None,replicas=2,rebalance=False):
		"""
		Create a new dyconf from arguments and write in a file

		Arguments:
			file -- path of the file to save
			hostnames -- names of the machines hosting the tokyo tyrant servers to use
			replicas -- number of replications for each key
			rebalance -- if True, if new parameters are detected, modify the mapping table
		"""
		self.needs_repair = False
		self.file = file
		# if the file already exists, load data from it
		if os.path.isfile(file):
			self.hostnames,self.replicas,self.slots,self.hosts_slots = self.load()
			if not hostnames:
				hostnames = self.hostnames
			# if the old parameters don't match the new ones, use rebalance if requested
			if set(hostnames) != set(self.hostnames) or replicas != self.replicas:
				if rebalance:
					self.rebalance(hostnames,replicas)
					self.needs_repair = True
				else:
					logging.error("New parameters are incompatible with old ones but rebalance wasn't requested... Set rebalance to True if needed")
					return
		
		# initialize new DyConf data
		else:
			self.hostnames = hostnames
			self.replicas = replicas
			self.new()

		# write data to file
		self.dump()

	def new(self):
		"""
		Create a new slot repartition for specified hostnames and replicas number
		"""
		self.slots = {}
		self.hosts_slots = {}
		self.add_replicas(self.replicas)

	""" 
	More homogenous slots repartition, but not finished yet
	def add_replicas(self,replicas):
		bytes = set(map(chr,range(256)))
		for r in replicas:
			used_slots = set([])
			for host in self.dyconf.hostnames:
				self.hosts_slots[host] = self.hosts_slots.get(host,[])
				past_slots = reduce(lambda x,y: x|y,self.hosts_slots[host])
				new_slots = set(random.choice(list(bytes - (used_slots | past_slots))))
				self.hosts_slots[host].append(set(random.choice(list(bytes-used_slots),len(bytes)/len(self.dyconf.hostnames))))
				for byte in self.hosts_slots[host][-1]:
					hosts = self.slots.get(byte,[])
					hosts.append(host)
					self.slots[byte] = hosts
	"""

	def build_from_sample(self,keys_sample):
		"""
		Build new mapping tables to match as closely as possible an existing key sample
		"""
		bytes = map(chr,range(256))
		weights = {}

		# count hosts that countain keys for each slots
		for key,hosts in keys_sample.iteritems():
			byte = self.hash(key)
			weights[byte] = weights.get(byte,{})
			for host in hosts:
				weights[byte][host] = weights[byte].get(host,0)+1
			
		self.slots = {}
		self.hosts_slots = {}
		# For each slot, keep the most associated hosts or select randomly if there isn't enough
		for byte in weights:
			hosts = weights[byte].keys()
			hosts.sort(lambda x,y:cmp(weights[byte][x],weights[byte][y]))
			self.slots[byte] = []
			for r in range(self.replicas):
				if r >= len(hosts):
					self.slots[byte].append(random.choice(list(tmp_hosts-set(self.slots[byte]))))
				else:
					self.slots[byte].append(hosts[r])
	
	def add_replicas(self,replicas):
		"""
		Add replicates hosts for each key slots.
		Simple random distribution for now.
		"""
		bytes = map(chr,range(256))
		for r in range(replicas):
			tmp_hosts = set(self.hostnames)
			tmp_nb = {}
			for i,byte in enumerate(bytes):
				self.slots[byte] = self.slots.get(byte,[])
				h = random.choice(list(tmp_hosts-set(self.slots[byte])))
				self.slots[byte].append(h)
	
	def rebalance(self,hostnames,replicas):
		"""
		Create a new slots repartition from an existing one
		"""

		# reduce the number of hosts
		if set(hostnames) != set(self.hostnames):
			raise NotImplementedError("Re-partitioning for new set of hosts is not implemented yet.")

		# reduce the number of replicas
		if replicas < self.replicas:
			for slot in self.slots:
				self.slots[slot] = self.slots[slot][:replicas]

		# increase the number of replicas
		if replicas > self.replicas:
			self.add_replicas(replicas - self.replicas)

		self.replicas = replicas

	def load(self):
		"""
		Get data from the file, return false if the file doesn't exist
		"""
		return pickle.load(open(self.file,'r'))

	def dump(self):
		"""
		Write data in the file
		"""
		pickle.dump([self.hostnames,self.replicas,self.slots,self.hosts_slots], open(self.file,'w'))

	def get_hosts(self,key):
		"""
		Return all hosts for a key
		"""
		return self.slots[self.hash(key)]

	def hash(self,key):
		"""
		Hash function for partitioning keys
		"""
		return hashlib.md5(key).digest()[0]

